# template2sql.py
# usage: 通过samples样本集，生成问题模板与样本库的对应json为你教案
# python ./src/vector/template2sql.py --samples ./data/samples.jsonl

import os,sys
# 添加项目根目录到路径
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)

import argparse
import json
import sys
from pathlib import Path
from tools.clean_question import clean_question


def main():
    parser = argparse.ArgumentParser(description="从 question 生成 template，并附加到 samples 输出为 template2sql 格式")
    parser.add_argument('--samples', type=str, required=True, help='输入的 samples JSONL 文件路径')

    args = parser.parse_args()

    input_path = Path(args.samples)
    output_path = Path('./tmp_template2sql.json')

    if not input_path.exists():
        print(f"Error: 输入文件不存在: {input_path}", file=sys.stderr)
        sys.exit(1)

    try:
        with input_path.open('r', encoding='utf-8') as fin, \
             output_path.open('w', encoding='utf-8') as fout:

            for line_num, line in enumerate(fin, 1):
                line = line.strip()
                if not line:
                    continue

                try:
                    sample = json.loads(line)
                except json.JSONDecodeError as e:
                    print(f"Error: 第 {line_num} 行 JSON 解析失败: {e}", file=sys.stderr)
                    sys.exit(1)

                if 'question' not in sample:
                    print(f"Warning: 第 {line_num} 行缺少 'question' 字段，跳过。", file=sys.stderr)
                    continue

                question = sample['question']
                try:
                    template = clean_question(question)
                except Exception as e:
                    print(f"Error: clean_question 处理问题 '{question}' 时出错: {e}", file=sys.stderr)
                    sys.exit(1)

                # 构造输出对象，将 template 插入到 question 前面
                output_obj = {
                    'id': sample['id'],
                    'template': template,
                    'question': sample['question'],
                }
                # 保留其他字段
                for k, v in sample.items():
                    if k not in output_obj:
                        output_obj[k] = v

                # 写入输出文件
                fout.write(json.dumps(output_obj, ensure_ascii=False) + '\n')

        print(f"Success: 已生成模板并保存至 {output_path}")

    except Exception as e:
        print(f"Error: 处理文件时发生错误: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()